"""
阿里云百炼平台 API获取文本向量
https://help.aliyun.com/zh/model-studio/developer-reference/text-embedding-synchronous-api?spm=a2c4g.11186623.0.i36

* text_type：
取值：query 或者 document，默认值为 document
说明：文本转换为向量后可以应用于检索、聚类、分类等下游任务，
对检索这类非对称任务为了达到更好的检索效果建议区分查询文本（query）和底库文本（document）类型,
聚类、分类等对称任务可以不用特殊指定，采用系统默认值"document"即可
"""

from typing import List, Union, Any, Optional
from openai import OpenAI
#"BAAI/bge-m3"


class EmbeddingModel:
    def __init__(self,
                 api_key: str,
                 base_url: str = "https://dashscope.aliyuncs.com/compatible-mode/v1",
                 model: str = 'text-embedding-v3',
                 dimension: int = 1024,
                 header: dict = None):

        self.model_name: str = model
        self.dimension: int = dimension
        self.api_key = api_key
        self.base_url = base_url
        self.headers = header or {
            'Authorization': self.api_key,
            'Content-Type': 'application/json'
        }

        # 客户端
        self.client = OpenAI(
            api_key=self.api_key,
            base_url=self.base_url,
            default_headers=self.headers
        )

    def get_text_embeddings(self,
                            texts: List[str]) -> Union[list[None], list[list]]:
        """
        获取文本的embedding
        :param texts: 文本列表，每次最多支持 25 条
        :return: 包含嵌入表示的响应结果
        """

        if not texts or not isinstance(texts, list):
            raise ValueError("texts 参数必须为非空的列表")

        embeddings = []

        try:
            completion = self.client.embeddings.create(
                model=self.model_name,
                input=texts,
                dimensions=self.dimension,
                encoding_format="float")

            for data in completion.data:
                embeddings.append(data.embedding)

            return embeddings

        except Exception as e:
            print(e)
            return [None] * len(texts)

    def get_text_embedding(self, text: str) -> Optional[List[float]]:
        embeddings = self.get_text_embeddings([text])
        return embeddings[0] if embeddings else None

    def get_embedding_dimension(self):
        """
        获取向量维度
        """
        return self.dimension

    def __getstate__(self):
        state = self.__dict__.copy()
        del state['client']
        return state

    def __setstate__(self, state):
        self.__dict__.update(state)
        self.client = OpenAI(api_key=self.api_key, base_url=self.base_url)


if __name__ == "__main__":
    apikey = "sk-68ac5f5ccf3540ba834deeeaecb48987"
    embd = EmbeddingModel(api_key=apikey)

    txts = [
        "风急天高猿啸哀",
        "渚清沙白鸟飞回",
        "无边落木萧萧下",
        "不尽长江滚滚来"
    ]

    result = embd.get_text_embeddings(
        txts
    )

    print(result)

